package main

import (
	"context"
	"fmt"
	"sync"
	"time"
)

func input(ctx context.Context, nums ...int) <-chan int {
	c := make(chan int)
	go func() {
		defer close(c)
		for _, n := range nums {
			select {
			case <-ctx.Done():
				fmt.Println("input done:", ctx.Err())
				return
			default:
				c <- n
				time.Sleep(time.Millisecond)
			}
		}
	}()
	return c
}

func square(ctx context.Context, in <-chan int) <-chan int {
	c := make(chan int)
	go func() {
		defer close(c)
		for n := range in {
			select {
			case <-ctx.Done():
				fmt.Println("square done:", ctx.Err())
				return
			default:
				c <- n * n
				time.Sleep(time.Second)
			}
		}
	}()
	return c
}

func merge(ctx context.Context, cs ...<-chan int) <-chan int {
	o := make(chan int)

	var wg sync.WaitGroup
	wg.Add(len(cs))

	for _, c := range cs {
		go func(c <-chan int) {
			defer wg.Done()
			for n := range c {
				select {
				case <-ctx.Done():
					return
				default:
					o <- n
					time.Sleep(time.Millisecond)
				}
			}
		}(c)
	}

	go func() {
		wg.Wait()
		close(o)
	}()

	return o
}

func main() {
	{
		ctx, done := context.WithCancel(context.Background())
		c := input(ctx, 1, 2, 3)
		fmt.Println(<-c)
		done()
		time.Sleep(time.Second)
		fmt.Println("example 1 done")
	}

	{
		ctx, done := context.WithCancel(context.Background())
		c := square(ctx, input(ctx, 1, 2, 3))
		fmt.Println(<-c)
		done()
		time.Sleep(time.Second)
		fmt.Println("example 2 done")
	}

	{
		ctx, done := context.WithCancel(context.Background())
		c := merge(ctx, square(ctx, input(ctx, 1)), square(ctx, input(ctx, 2)))
		fmt.Println(<-c)
		done()
		time.Sleep(time.Second)
		fmt.Println("example 3 done")
	}

	{
		ctx1 := context.WithValue(context.Background(), "language", "Go")
		ctx2, _ := context.WithTimeout(ctx1, 2*time.Second)
		time.Sleep(1500 * time.Millisecond)
		for n := range square(ctx2, input(ctx2, 1, 2, 3)) {
			fmt.Println(n)
		}

		fmt.Println(ctx1.Value("language"))
		fmt.Println(ctx2.Value("language"))
		time.Sleep(time.Second)
		fmt.Println("example 4 done")
	}
	{
		wait := sync.WaitGroup{}
		ctxParent, cancelParent := context.WithCancel(context.Background())
		wait.Add(1)
		go func() {
			defer wait.Done()
			select {
			case <-ctxParent.Done():
				fmt.Println("parent done with error:", ctxParent.Err())
			case <-time.After(time.Second):
				fmt.Println("parent done failed")
			}
		}()

		ctxChild, _ := context.WithCancel(ctxParent)
		wait.Add(1)
		go func() {
			defer wait.Done()
			select {
			case <-ctxChild.Done():
				fmt.Println("child done with error:", ctxParent.Err())
			case <-time.After(time.Second):
				fmt.Println("child done failed")
			}
		}()

		cancelParent() // parent will clean child context
		wait.Wait()
	}

	main2()
}
